from typing import Optional
import tkinter as tk
from tkinter import filedialog

import cv2
import numpy as np
from numpy import ndarray as mat


class RGBSelectorTool:
    def __init__(self, windows_width: int = 720, windows_height: int = 640, windows_name: str = "RGB Selector Tool"):
        self.__windows_width: int = windows_width
        self.__windows_height: int = windows_height
        self.__windows_name: str = windows_name
        self.__img_rgb: Optional[mat] = None

    def run_bgr(self, img_bgr):
        # 将BGR转换为RGB（OpenCV默认是BGR格式）
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        self.run_rgb(img_rgb)

    def run_rgb(self, img_rgb):
        self.__img_rgb = img_rgb
        cv2.namedWindow(self.__windows_name, cv2.WINDOW_NORMAL)
        debug_resize_times = 0.6
        cv2.resizeWindow(self.__windows_name, np.int_(self.__windows_width * debug_resize_times),
                         np.int_(self.__windows_height * debug_resize_times * 1.5))
        
        # 创建RGB范围的滑动条（0-255）
        cv2.createTrackbar('R_Min', self.__windows_name, 0, 255, self.__cv_select_callback)
        cv2.createTrackbar('G_Min', self.__windows_name, 0, 255, self.__cv_select_callback)
        cv2.createTrackbar('B_Min', self.__windows_name, 0, 255, self.__cv_select_callback)
        cv2.createTrackbar('R_Max', self.__windows_name, 255, 255, self.__cv_select_callback)
        cv2.createTrackbar('G_Max', self.__windows_name, 255, 255, self.__cv_select_callback)
        cv2.createTrackbar('B_Max', self.__windows_name, 255, 255, self.__cv_select_callback)

        min_r, min_g, min_b, max_r, max_g, max_b = 0, 0, 0, 255, 255, 255

        while True:
            min_r = cv2.getTrackbarPos('R_Min', self.__windows_name)
            min_g = cv2.getTrackbarPos('G_Min', self.__windows_name)
            min_b = cv2.getTrackbarPos('B_Min', self.__windows_name)
            max_r = cv2.getTrackbarPos('R_Max', self.__windows_name)
            max_g = cv2.getTrackbarPos('G_Max', self.__windows_name)
            max_b = cv2.getTrackbarPos('B_Max', self.__windows_name)

            range_min = np.array([min_r, min_g, min_b])
            range_max = np.array([max_r, max_g, max_b])
            print(f"Scalar lower_rgb({min_r}, {min_g}, {min_b});")
            print(f"Scalar upper_rgb({max_r}, {max_g}, {max_b});")

            # 在RGB颜色空间中进行范围筛选
            img_rgb_debug = cv2.inRange(self.__img_rgb, range_min, range_max)
            
            # 创建彩色显示图像：将筛选结果叠加到原图上
            img_masked = self.__img_rgb.copy()
            img_masked[img_rgb_debug == 0] = [0, 0, 0]  # 不在范围内的像素设为黑色
            
            # 转换回BGR格式用于OpenCV显示
            img_display = cv2.cvtColor(img_masked, cv2.COLOR_RGB2BGR)
            
            # 垂直拼接原图（转为BGR）、掩膜图和筛选结果图
            img_original_bgr = cv2.cvtColor(self.__img_rgb, cv2.COLOR_RGB2BGR)
            img_mask_colored = cv2.applyColorMap(img_rgb_debug, cv2.COLORMAP_JET)
            
            # 调整图像大小以便显示
            height = img_original_bgr.shape[0] // 3
            img_original_resized = cv2.resize(img_original_bgr, (0, 0), fx=1, fy=height/img_original_bgr.shape[0])
            img_mask_resized = cv2.resize(img_mask_colored, (0, 0), fx=1, fy=height/img_mask_colored.shape[0])
            img_result_resized = cv2.resize(img_display, (0, 0), fx=1, fy=height/img_display.shape[0])
            
            # 垂直拼接三个图像
            img_combined = np.vstack([img_original_resized, img_mask_resized, img_result_resized])
            
            cv2.imshow(self.__windows_name, img_combined)

            key = cv2.waitKey(10)
            if key & 0xFF == ord('q'):
                cv2.destroyWindow(self.__windows_name)
                break
            elif key & 0xFF == ord('s'):
                # 按's'键保存当前参数到文件
                self.__save_parameters(min_r, min_g, min_b, max_r, max_g, max_b)

        return min_r, min_g, min_b, max_r, max_g, max_b

    def __save_parameters(self, min_r, min_g, min_b, max_r, max_g, max_b):
        """保存RGB参数到文件"""
        with open('rgb_parameters.txt', 'w') as f:
            f.write(f"RGB Range Parameters:\n")
            f.write(f"R_MIN: {min_r}\n")
            f.write(f"G_MIN: {min_g}\n")
            f.write(f"B_MIN: {min_b}\n")
            f.write(f"R_MAX: {max_r}\n")
            f.write(f"G_MAX: {max_g}\n")
            f.write(f"B_MAX: {max_b}\n")
            f.write(f"\nNumPy arrays:\n")
            f.write(f"range_min = np.array([{min_r}, {min_g}, {min_b}])\n")
            f.write(f"range_max = np.array([{max_r}, {max_g}, {max_b}])\n")
        print(f"Parameters saved to rgb_parameters.txt")

    @staticmethod
    def __cv_select_callback(value: int):
        pass


class RGBSelectorToolAdvanced(RGBSelectorTool):
    """高级RGB调参助手，增加了更多功能"""
    
    def run_rgb(self, img_rgb):
        self.__img_rgb = img_rgb
        cv2.namedWindow(self.__windows_name, cv2.WINDOW_NORMAL)
        debug_resize_times = 0.8
        cv2.resizeWindow(self.__windows_name, np.int_(self.__windows_width * debug_resize_times),
                         np.int_(self.__windows_height * debug_resize_times * 1.2))
        
        # 创建RGB范围的滑动条
        cv2.createTrackbar('R_Min', self.__windows_name, 0, 255, self.__cv_select_callback)
        cv2.createTrackbar('G_Min', self.__windows_name, 0, 255, self.__cv_select_callback)
        cv2.createTrackbar('B_Min', self.__windows_name, 0, 255, self.__cv_select_callback)
        cv2.createTrackbar('R_Max', self.__windows_name, 255, 255, self.__cv_select_callback)
        cv2.createTrackbar('G_Max', self.__windows_name, 255, 255, self.__cv_select_callback)
        cv2.createTrackbar('B_Max', self.__windows_name, 255, 255, self.__cv_select_callback)
        
        # 添加显示模式选择
        cv2.createTrackbar('Mode', self.__windows_name, 0, 2, self.__cv_select_callback)  # 0:掩膜, 1:彩色结果, 2:组合显示

        while True:
            min_r = cv2.getTrackbarPos('R_Min', self.__windows_name)
            min_g = cv2.getTrackbarPos('G_Min', self.__windows_name)
            min_b = cv2.getTrackbarPos('B_Min', self.__windows_name)
            max_r = cv2.getTrackbarPos('R_Max', self.__windows_name)
            max_g = cv2.getTrackbarPos('G_Max', self.__windows_name)
            max_b = cv2.getTrackbarPos('B_Max', self.__windows_name)
            mode = cv2.getTrackbarPos('Mode', self.__windows_name)

            range_min = np.array([min_r, min_g, min_b])
            range_max = np.array([max_r, max_g, max_b])
            
            # 实时显示像素统计信息
            img_rgb_debug = cv2.inRange(self.__img_rgb, range_min, range_max)
            pixel_count = np.sum(img_rgb_debug == 255)
            total_pixels = img_rgb_debug.shape[0] * img_rgb_debug.shape[1]
            percentage = (pixel_count / total_pixels) * 100
            
            print(f"Range [R, G, B]: MIN - [{min_r}, {min_g}, {min_b}], MAX - [{max_r}, {max_g}, {max_b}] | "
                  f"Pixels: {pixel_count}/{total_pixels} ({percentage:.2f}%)")

            # 根据模式显示不同的结果
            if mode == 0:  # 仅显示掩膜
                cv2.imshow(self.__windows_name, img_rgb_debug)
            elif mode == 1:  # 显示彩色筛选结果
                img_masked = self.__img_rgb.copy()
                img_masked[img_rgb_debug == 0] = [0, 0, 0]
                img_display = cv2.cvtColor(img_masked, cv2.COLOR_RGB2BGR)
                cv2.imshow(self.__windows_name, img_display)
            else:  # 组合显示（原图、掩膜、结果）
                img_original_bgr = cv2.cvtColor(self.__img_rgb, cv2.COLOR_RGB2BGR)
                img_masked = self.__img_rgb.copy()
                img_masked[img_rgb_debug == 0] = [0, 0, 0]
                img_result = cv2.cvtColor(img_masked, cv2.COLOR_RGB2BGR)
                
                # 水平拼接
                img_combined = np.hstack([
                    cv2.resize(img_original_bgr, (240, 180)),
                    cv2.resize(cv2.cvtColor(img_rgb_debug, cv2.COLOR_GRAY2BGR), (240, 180)),
                    cv2.resize(img_result, (240, 180))
                ])
                cv2.imshow(self.__windows_name, img_combined)

            key = cv2.waitKey(10)
            if key & 0xFF == ord('q'):
                cv2.destroyWindow(self.__windows_name)
                break
            elif key & 0xFF == ord('s'):
                self.__save_parameters(min_r, min_g, min_b, max_r, max_g, max_b)

        return min_r, min_g, min_b, max_r, max_g, max_b


if __name__ == '__main__':
    root = tk.Tk()
    root.withdraw()
    file_path = filedialog.askopenfilename()
    print("Selected file:", file_path)
    
    # 选择使用基础版本还是高级版本
    # use_advanced = input("Use advanced version? (y/n): ").lower() == 'y'
    use_advanced = False
    
    if use_advanced:
        tool = RGBSelectorToolAdvanced()
        print("Advanced RGB Selector Tool")
        print("Mode 0: Mask only")
        print("Mode 1: Colored result")
        print("Mode 2: Combined view")
        print("Press 's' to save parameters, 'q' to quit")
    else:
        tool = RGBSelectorTool()
        print("Basic RGB Selector Tool")
        print("Press 's' to save parameters, 'q' to quit")
    
    result = tool.run_bgr(cv2.imread(file_path))
    print(f"Final RGB range: MIN [{result[0]}, {result[1]}, {result[2]}], MAX [{result[3]}, {result[4]}, {result[5]}]")